/*
 * bst/avl_tree.hh - AVL tree implementation.
 * Copyright (C) 2008  Mikhail Vorozhtsov
 * See the LICENSE section of the README file for details.
 */

/* $Id$ */

#ifndef HEADER_LIBBST_MJ_AVL_TREE_HH
#define HEADER_LIBBST_MJ_AVL_TREE_HH

#include HEADER (LIBBST_PREFIX,base.hh)

START_NAMESPACE (LIBBST_NS)

enum avl_balance_t
{
  AVL_BALANCE_LEFT = 0,
  AVL_BALANCE_OK = 1,
  AVL_BALANCE_RIGHT = 2
};

template <typename P>
struct basic_avl_node_t : public P
{
private:
  typedef basic_avl_node_t node_t;

public:
  INLINE
  node_t *
  get_left () const
  {
    return (node_t *) (this->left.raw & P::BASE_MASK);
  }

  INLINE
  void
  set_left (node_t *x)
  {
    this->left.base = (u64_t) x >> P::BASE_SHIFT;
  }

  INLINE
  avl_balance_t
  get_balance () const
  {
    return (avl_balance_t) this->left.bits12;
  }

  INLINE
  void
  set_balance (avl_balance_t x)
  {
    this->left.bits12 = x;
  }
};

template <typename NODE, typename TREE, typename GLUE>
struct fwd_avl_tree_t : public base_tree_t<NODE, TREE, GLUE>
{
public:
  template <typename STACK>
  bool
  insert (typename NODE::container_t *x, STACK &stack)
  {
    NODE *parent, *prev_or_next;
    bool left_p;

    if (full_lookup (x, parent, left_p, prev_or_next, stack))
      return false;

    NODE *n = (NODE *) NODE::from_container (x);
    before_insert (n, parent, left_p);
    n->insert (this, parent, left_p, prev_or_next);
    n->set_balance (AVL_BALANCE_OK);

    if (parent)
      stack.pop ();

    while (n != this->get_root ())
      {
        if (parent == NULL)
          parent = stack.pop (this, n);

        avl_balance_t bal = parent->get_balance ();

        if (n == parent->get_left ())
          {
            if (bal == AVL_BALANCE_RIGHT)
              {
                parent->set_balance (AVL_BALANCE_OK);
                break;
              }

            if (bal == AVL_BALANCE_OK)
              {
                parent->set_balance (AVL_BALANCE_LEFT);
                n = parent;
                parent = NULL;
                continue;
              }

            NODE *grandparent = parent == this->get_root () ? NULL
                                : stack.pop (this, parent);

            if (n->get_balance () == AVL_BALANCE_LEFT)
              {
                parent->rotate_right (this, grandparent);
                n->set_balance (AVL_BALANCE_OK);
                parent->set_balance (AVL_BALANCE_OK);
              }
            else
              {
                NODE *child = n->get_right ();
                bal = child->get_balance ();

                n->rotate_left (this, parent);
                parent->rotate_right (this, grandparent);

                if (bal == AVL_BALANCE_LEFT)
                  {
                    n->set_balance (AVL_BALANCE_OK);
                    parent->set_balance (AVL_BALANCE_RIGHT);
                  }
                else if (bal == AVL_BALANCE_RIGHT)
                  {
                    n->set_balance (AVL_BALANCE_LEFT);
                    parent->set_balance (AVL_BALANCE_OK);
                  }
                else
                  {
                    n->set_balance (AVL_BALANCE_OK);
                    parent->set_balance (AVL_BALANCE_OK);
                  }

                child->set_balance (AVL_BALANCE_OK);
              }

            break;
          }
        else
          {
            if (bal == AVL_BALANCE_LEFT)
              {
                parent->set_balance (AVL_BALANCE_OK);
                break;
              }

            if (bal == AVL_BALANCE_OK)
              {
                parent->set_balance (AVL_BALANCE_RIGHT);
                n = parent;
                parent = NULL;
                continue;
              }

            NODE *grandparent = parent == this->get_root () ? NULL
                                : stack.pop (this, parent);

            if (n->get_balance () == AVL_BALANCE_RIGHT)
              {
                parent->rotate_left (this, grandparent);
                n->set_balance (AVL_BALANCE_OK);
                parent->set_balance (AVL_BALANCE_OK);
              }
            else
              {
                NODE *child = n->get_left ();
                bal = child->get_balance ();

                n->rotate_right (this, parent);
                parent->rotate_left (this, grandparent);

                if (bal == AVL_BALANCE_RIGHT)
                  {
                    n->set_balance (AVL_BALANCE_OK);
                    parent->set_balance (AVL_BALANCE_LEFT);
                  }
                else if (bal == AVL_BALANCE_LEFT)
                  {
                    n->set_balance (AVL_BALANCE_RIGHT);
                    parent->set_balance (AVL_BALANCE_OK);
                  }
                else
                  {
                    n->set_balance (AVL_BALANCE_OK);
                    parent->set_balance (AVL_BALANCE_OK);
                  }

                child->set_balance (AVL_BALANCE_OK);
              }

            break;
          }
      }

    return true;
  }

  template <typename KEY, typename CMP, typename STACK>
  typename NODE::container_t *
  remove_key (const KEY &key, STACK &stack)
  {
    NODE *parent, *prev_or_next;
    bool left_p;
    typename NODE::container_t *x
      = this->template full_lookup_key<KEY, CMP, STACK>
                         (key, parent, left_p, prev_or_next, stack);

    if (!x)
      return NULL;

    NODE *grandparent = NULL;
    NODE *n = (NODE *) NODE::from_container (x);
    NODE *left = (NODE *) n->get_left ();
    NODE *right = (NODE *) n->get_right ();
    NODE *child;

    if (left)
      {
        NODE *pred_parent = n;
        NODE *pred = left->get_max_and_replace (pred_parent, stack);
        NODE *succ = NULL;
        child = n->replace_with_pred (this, parent, left_p, pred, pred_parent,
                                      succ);

        pred->set_balance (n->get_balance ());
        after_remove (n, parent, pred, right != NULL, succ);

        if (pred_parent == n)
          {
            grandparent = parent;
            parent = pred;
            left_p = true;

            if (grandparent)
              stack.pop ();
          }
        else
          {
            parent = pred_parent;
            left_p = false;
          }
      }
    else if (right)
      {
        NODE *succ_parent = n;
        NODE *succ = right->get_min_and_replace (succ_parent, stack);
        child = n->replace_with_succ (this, parent, left_p, succ, succ_parent,
                                      prev_or_next);

        succ->set_balance (n->get_balance ());
        after_remove (n, parent, NULL, true, succ);

        if (succ_parent == n)
          {
            grandparent = parent;
            parent = succ;
            left_p = false;

            if (grandparent)
              stack.pop ();
          }
        else
          {
            parent = succ_parent;
            left_p = true;
          }
      }
    else
      {
        n->remove (this, parent, left_p, prev_or_next);
        after_remove (n, parent, NULL, false, NULL);
        child = NULL;
      }

    if (parent)
      stack.pop ();

    for (n = child; n != this->get_root ();)
      {
        if (parent == NULL)
          {
            parent = stack.pop (this, n);
            left_p = n == parent->get_left ();
          }

        avl_balance_t bal = parent->get_balance ();

        if (left_p)
          {
            if (bal == AVL_BALANCE_OK)
              {
                parent->set_balance (AVL_BALANCE_RIGHT);
                break;
              }

            if (grandparent == NULL)
              {
                if (parent == this->get_root ())
                  grandparent = NULL;
                else
                  {
                    grandparent = stack.pop (this, parent);
                    left_p = parent == grandparent->get_left ();
                  }
              }
            else
              left_p = parent == grandparent->get_left ();

            if (bal == AVL_BALANCE_LEFT)
              {
                parent->set_balance (AVL_BALANCE_OK);
                n = parent;
                parent = grandparent;
                grandparent = NULL;
                continue;
              }

            NODE *sibling = parent->get_right ();
            bal = sibling->get_balance ();

            if (bal == AVL_BALANCE_LEFT)
              {
                NODE *nephew = sibling->get_left (); 
                bal = nephew->get_balance ();

                sibling->rotate_right (this, parent);
                parent->rotate_left (this, grandparent);

                if (bal == AVL_BALANCE_LEFT)
                  {
                    parent->set_balance (AVL_BALANCE_OK);
                    sibling->set_balance (AVL_BALANCE_RIGHT);
                    nephew->set_balance (AVL_BALANCE_OK);
                  }
                else if (bal == AVL_BALANCE_OK)
                  {
                    parent->set_balance (AVL_BALANCE_OK);
                    sibling->set_balance (AVL_BALANCE_OK);
                  }
                else
                  {
                    parent->set_balance (AVL_BALANCE_LEFT);
                    sibling->set_balance (AVL_BALANCE_OK);
                    nephew->set_balance (AVL_BALANCE_OK);
                  }

                n = nephew;
                parent = grandparent;
                grandparent = NULL;
              }
            else
              {
                parent->rotate_left (this, grandparent);

                if (bal == AVL_BALANCE_OK)
                  {
                    sibling->set_balance (AVL_BALANCE_LEFT);
                    parent->set_balance (AVL_BALANCE_RIGHT);
                    break;
                  }

                sibling->set_balance (AVL_BALANCE_OK);
                parent->set_balance (AVL_BALANCE_OK);
                n = sibling;
                parent = grandparent;
                grandparent = NULL;
              }
          }
        else
          {
            if (bal == AVL_BALANCE_OK)
              {
                parent->set_balance (AVL_BALANCE_LEFT);
                break;
              }

            if (grandparent == NULL)
              {
                if (parent == this->get_root ())
                  grandparent = NULL;
                else
                  {
                    grandparent = stack.pop (this, parent);
                    left_p = parent == grandparent->get_left ();
                  }
              }
            else
              left_p = parent == grandparent->get_left ();

            if (bal == AVL_BALANCE_RIGHT)
              {
                parent->set_balance (AVL_BALANCE_OK);
                n = parent;
                parent = grandparent;
                grandparent = NULL;
                continue;
              }

            NODE *sibling = parent->get_left ();
            bal = sibling->get_balance ();

            if (bal == AVL_BALANCE_RIGHT)
              {
                NODE *nephew = sibling->get_right (); 
                bal = nephew->get_balance ();

                sibling->rotate_left (this, parent);
                parent->rotate_right (this, grandparent);

                if (bal == AVL_BALANCE_LEFT)
                  {
                    parent->set_balance (AVL_BALANCE_RIGHT);
                    sibling->set_balance (AVL_BALANCE_OK);
                    nephew->set_balance (AVL_BALANCE_OK);
                  }
                else if (bal == AVL_BALANCE_OK)
                  {
                    parent->set_balance (AVL_BALANCE_OK);
                    sibling->set_balance (AVL_BALANCE_OK);
                  }
                else
                  {
                    parent->set_balance (AVL_BALANCE_OK);
                    sibling->set_balance (AVL_BALANCE_LEFT);
                    nephew->set_balance (AVL_BALANCE_OK);
                  }

                n = nephew;
                parent = grandparent;
                grandparent = NULL;
              }
            else
              {
                parent->rotate_right (this, grandparent);

                if (bal == AVL_BALANCE_OK)
                  {
                    sibling->set_balance (AVL_BALANCE_RIGHT);
                    parent->set_balance (AVL_BALANCE_LEFT);
                    break;
                  }

                sibling->set_balance (AVL_BALANCE_OK);
                parent->set_balance (AVL_BALANCE_OK);
                n = sibling;
                parent = grandparent;
                grandparent = NULL;
              }
          }
      }

    return x;
  }

public:
  void
  replace (typename NODE::container_t *x, typename NODE::container_t *with)
  {
    NODE *node = NODE::from_container (x);
    NODE *with_node = NODE::from_container (with);

    with_node->set_balance (node->get_balance ());
    base_tree_t<NODE, TREE, GLUE>::replace (x, with);
  }
};

template <typename NODE, typename TREE = basic_tree_t, typename GLUE = glue_t>
struct avl_tree_t
  : public wrap_tree_t<ord_lookup_tree_t<fwd_avl_tree_t<NODE, TREE, GLUE> > >
{
};

END_NAMESPACE

#endif /* HEADER_LIBBST_MJ_AVL_TREE_HH */

